{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "819ff1f1-2c7e-484b-8697-846a84be9504",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 知识库的构建\n",
    "### 文档解析-->分块-->embedding-->向量数据库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "27135fa1-04fa-48a1-a611-bc26af5ab84e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import RagEmbedding, RagLLM\n",
    "from doc_parse import chunk, read_and_process_excel, logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fc02ea7a-97ab-46fb-b742-60e73f5bdea9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "import chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f2580ace-b5e7-46a0-b10a-2ede2e7a46e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf_files = ['./data/zhidu_employee.pdf',\n",
    "            './data/zhidu_travel.pdf']\n",
    "\n",
    "excel_files = ['./data/zhidu_detail.xlsx']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "30122d67-a7a1-4ace-ba3a-1730fac10c23",
   "metadata": {},
   "outputs": [],
   "source": [
    "r_spliter = RecursiveCharacterTextSplitter(chunk_size=128,\n",
    "                                          chunk_overlap=30,\n",
    "                                          separators=[\"\\n\\n\",\n",
    "                                                      \"\\n\", \n",
    "                                                      \".\",\n",
    "                                                      \"\\uff0e\",  # Fullwidth full stop\n",
    "                                                       \"\\u3002\",  # Ideographic full stop\n",
    "                                                      \",\",\n",
    "                                                      \"\\uff0c\",  # Fullwidth comma\n",
    "                                                      \"\\u3001\",  # Ideographic comma\n",
    "                                                     ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "61d84fb3-9551-4c2e-b47c-38f555db91aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OCR is running...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Miss outlines\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OCR finished.\n",
      "OCR: 1.9347431138157845\n",
      "preprocess\n",
      "Layout analysis finished.\n",
      "layouts: 2.390368826687336\n",
      "preprocess\n",
      "Table analysis finished.\n",
      "Text merging finished\n",
      "OCR is running...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Miss outlines\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OCR finished.\n",
      "OCR: 0.9043905958533287\n",
      "preprocess\n",
      "Layout analysis finished.\n",
      "layouts: 1.2087930291891098\n",
      "preprocess\n",
      "Table analysis finished.\n",
      "Text merging finished\n"
     ]
    }
   ],
   "source": [
    "doc_data = []\n",
    "for pdf_file_name in pdf_files:\n",
    "    res = chunk(pdf_file_name, callback=logger)\n",
    "    for data in res:\n",
    "        content = data[\"content_with_weight\"]\n",
    "        if '<table>' not in content and len(content) > 200:\n",
    "            doc_data = doc_data + r_spliter.split_text(content)\n",
    "        else:\n",
    "            doc_data.append(content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7956fd4b-8ef3-48ec-92a7-af270e9c9b57",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "530 ========== <table><caption>病假发放标准：</caption>\n",
      "<tr><td  >病假时间</td><td  >连续工龄</td><td  >发放标准</td></tr>\n",
      "<tr><td></td><td  >不满二年</td><td  >60% </td></tr>\n",
      "<tr><td></td><td  >已满二年不满四年</td><td  >70% </td></tr>\n",
      "<tr><td  >6 个月以内病假</td><td  >已满四年不满六年</td><td  >80% </td></tr>\n",
      "<tr><td></td><td  >已满六年不满八年</td><td  >90% </td></tr>\n",
      "<tr><td></td><td  >八年以上</td><td  >100% </td></tr>\n",
      "<tr><td></td><td  >不满一年</td><td  >40% </td></tr>\n",
      "<tr><td  >6 个月以上病假</td><td  >已满一年不满三年</td><td  >50% </td></tr>\n",
      "<tr><td></td><td  >连续工龄三年以上</td><td  >60% </td></tr>\n",
      "</table>\n",
      "58 ========== 教职工考勤管理制度\n",
      "第一节适用范围\n",
      "1、本制度包括了考勤、休假、加班等方面的规定。2、本制度适用于学校全体教职员工。\n",
      "7 ========== 第二节考勤规定\n",
      "66 ========== 1、学校的工作时间由学校决定并公布。学校内除特殊岗位特别规定外，全体教职员工均应严格执行学校的作息时间表，不迟到、不早退、不中途离校\n",
      "120 ========== 。工作时间：星期一至星期四7:55-16:55 星期五7:55-16:152、所有教职工实行考勤打卡制度，工作日内，每天需打卡两次，上午上班一次和下午下班一次。3、教职员工因故(特殊情况除外)晚到或早退，应事先履行请假手续，经批准后方可离校\n",
      "59 ========== 。原则上，员工请假无论时间长短、假期形式，除急诊病假或突发事件外，一律需按照请假流程。请假，需事先在钉钉系统中提交申请\n",
      "102 ========== 。请假，需事先在钉钉系统中提交申请。有效的请假流程为：（1）员工休假必须事先向部门负责人申请，将工作交接清楚方可休假；（2）2 天以内的假期必须经过部门分管领导审批；（3）3 天以上的假期必须经过校长审批\n",
      "125 ========== 。如遇紧急情况，口头申请请假的，应在上班后两天内办理补请假手续，未在规定时间内办理的，逾期无效，按旷工处理。4、无工作理由，超过上班时间到岗的，视为迟到；未到下班时间提前离校的，视为早退；中途未经批准离校，视为旷工；迟到、早退、旷工者按照相关办法处理\n",
      "44 ========== 。5、考勤员徇私舞弊、弄虚作假的，按学校奖惩规定处理。情节严重的，予以通报批评直至解聘。\n",
      "1046 ========== 第三节休假规定\n",
      "休假分以下八种：事假、病假、婚假、丧假、产假、哺乳假、工伤假、调休。1、事假：（1）请假必须严格按照学校规定的请假程序，应由本人以书面形式，应在前两日办理手续，特殊情况应通过电话或者口头请假，应在事假结束于2 个工作日内完成补办请假手续。未经请假或请假未准而擅自离岗者，以旷工论处。（2）事假最小计算单位为半天，事假一次不得超过3 天。（3）事假：基本工资和岗位津贴均按请假天数占实际上班天数比例来算。（4）请假理由不充分或致工作妨碍时，可酌情缩短假期、或延期给假、或不予给假。（5）请假者必须将所任课务或经办事务交待给代理人员，并于请假单内注明。\n",
      "2、病假：因身体健康问题不能正常工作的员工可申请病假，休假后须提供三级医院开具的病假条或诊断证明。（1）教职工休病假需提前申请。如因情况紧急或突发情况无法请假的，应通过电话或者口头请假，应在病假结束于2 个工作日内补办相关手续。未经请假或请假未准而擅自离岗者，以旷工论处。(2)病假按照工龄系数，对病假日工资进行扣除。\n",
      "<table><caption>病假发放标准：</caption>\n",
      "<tr><td  >病假时间</td><td  >连续工龄</td><td  >发放标准</td></tr>\n",
      "<tr><td></td><td  >不满二年</td><td  >60% </td></tr>\n",
      "<tr><td></td><td  >已满二年不满四年</td><td  >70% </td></tr>\n",
      "<tr><td  >6 个月以内病假</td><td  >已满四年不满六年</td><td  >80% </td></tr>\n",
      "<tr><td></td><td  >已满六年不满八年</td><td  >90% </td></tr>\n",
      "<tr><td></td><td  >八年以上</td><td  >100% </td></tr>\n",
      "<tr><td></td><td  >不满一年</td><td  >40% </td></tr>\n",
      "<tr><td  >6 个月以上病假</td><td  >已满一年不满三年</td><td  >50% </td></tr>\n",
      "<tr><td></td><td  >连续工龄三年以上</td><td  >60% </td></tr>\n",
      "</table>\n",
      "3、婚假：全体教职工达到国家法律（男方22 周岁，女方20 周岁）规定的，经过法定程序办理结婚手续的，可凭结婚证申请婚假15 天（均含休息日）。\n",
      "158 ========== （1）婚假期间员工工资照发。\n",
      "（2）为了保证教学活动的顺利进行，提倡婚假安排在寒暑假。\n",
      "4、丧假：员工因配偶或直系亲属丧亡时，可申请丧假：配偶、父母（公婆、岳父母）、子女及其配偶、祖父母、外祖父母、孙子女（外孙子女）及其配偶、曾祖父母、曾外祖父母丧亡，丧假3 天。\n",
      "5、产假：符合计划生育政策规定的女性员工可享受产假\n",
      "113 ========== （1）女教职员工的法定产假为158 天，其中产前可休假15 天；\n",
      "（2）怀孕不满四个月流产的，可给予l5-30 天的产假；\n",
      "（3）怀孕满四个月以上流产的，可给予40 天的产假；\n",
      "（4）男教职员工在配偶分娩时可休15 天陪护假；\n",
      "113 ========== （4）男教职员工在配偶分娩时可休15 天陪护假；\n",
      "（5）女职工产前检查，可按照医院的日期，调整好课务之后去检查。每次产检准假半天，并提供产检相关凭证；\n",
      "（6）以上产假天数为日历天数，包含法定节假日和休息日。\n",
      "（7）产假期间薪资\n",
      "123 ========== （7）产假期间薪资\n",
      "根据《福建新计生条例第四十一条规定》：妇女怀孕、生育和哺乳期间，按照国家和本省有关规定享受特殊劳动保护并可以获得帮助和补偿。产假期间工资照发，不影响晋升。①保胎假，工资按照病假发\n",
      "保胎假需由医生开证明，所以按病假待遇发放工资。\n",
      "80 ========== 保胎假需由医生开证明，所以按病假待遇发放工资。\n",
      "②产前假，工资按照病假计算。\n",
      "③产假工资：产假期间领取生育津贴，生育津贴没有覆盖的产假工资按照基本工资标准发放。\n",
      "109 ========== 6、哺乳假：哺乳期行政部门、教辅岗位职工可在不影响工作的前提下，婴儿不满一周岁的，每天准假两个小时，上午一小时，下午一小时，但须事先在人事部备案。哺乳期教师岗位在保证教学工作正常开展、保证教学质量的前提下，可以不坐班。\n",
      "85 ========== 7、工伤假：员工在上下班途中遭遇机动车交通事故的（交管部门出具机动车事故报告），或工作时、因公外出执行工作任务期间，遭受身体伤害需要医护治疗的，根据国家有关规定给予工伤假\n",
      "127 ========== 。8、调休：各职能部门因工作需要，需要在工作日、节假日安排本部门加班或值班的，应由行政部门书面报送至人力资源部，如遇突发情况可事后补办手续。审批通过后，加班或值班时间可申请调休，教职工本人需填写《员工请假申请单》，报相关部门人员签字，并交给人事部核算考勤\n",
      "1 ========== 。\n",
      "83 ========== 9、以下情况视同出勤：持有学校书面通知，经部门负责人同意参加家长会的。以上，依本规则所请各假如发现有虚假行为者，一经查明，除以旷工论处外，并根据学校奖惩规定予以处罚。\n",
      "10 ========== 第四节加班和调休规定\n",
      "96 ========== 1、学校以下列日期为例行公休日(若有变更需事先公布)，但因学校工作需要可指定照常上班，以加班计算：(1)法定节假日(2)周六、周日2、正常工作日，延长工作时间连续达到4 小时以上的方可计算加班\n",
      "118 ========== 。（加班需提前申请），加班需有打卡记录，无打卡记录支撑的加班时间，不计加班。加班费按照实际加班时长的2倍计算。3、加班以调休等额返还（代替），凡调休人员应填写《请假申请表》，选择“调休”一栏，经所在部门分管领导签字后，交由人事处核实备案\n",
      "25 ========== 。4、行政岗及教辅岗都需要参与法定节假日轮流值班。\n",
      "125 ========== 第五节违规违纪处理\n",
      "1. 迟到、早退(1)劳动考勤是公司支付薪资的依据和职工年度岗位考核内容之一。（2）迟到或早退60分钟以上（含60 分钟），每次视同缺勤1 天。（3）职工迟到、早退、脱岗累计超过3 次的（含），从第1 次起，每次扣减工资50 元。\n",
      "119 ========== 2. 缺打卡\n",
      "员工上下班无打卡或者有效签注，又无有效证明的视为旷工或早退半天。\n",
      "3. 旷工\n",
      "员工有以下情形之一未能出勤的，视为旷工：\n",
      "⑴未按规定程序办理请假手续的或虽请假但未经批准，离开工作岗位不上班的；\n",
      "⑵提供虚假的证明材料获得准假的；\n",
      "102 ========== ⑵提供虚假的证明材料获得准假的；\n",
      "⑶迟到、早退或擅离学校达1 小时以上的；\n",
      "⑷谎报请假原因，骗取、伪造、涂改病例证明或领导签字的；\n",
      "⑸事假未经批准的；\n",
      "⑹假满未经续假或未被同意续假的，或无出差审批手续的；\n",
      "97 ========== ⑹假满未经续假或未被同意续假的，或无出差审批手续的；\n",
      "⑺不服从学校应工作需要进行的正常调动和分配，有岗不上的；\n",
      "⑻以公出或病、事假为由，工作时间在外干私活的；⑼其它法律规定的情形。\n",
      "4.违规处理\n",
      "25 ========== (1)无故缺勤、未经批准备案的假期均按照旷工处理；\n",
      "63 ========== (2)旷工以0.5 天为单位，按照缺勤时间的三倍扣发工资；(3)连续旷工超过3 天，或全年累计旷工超过7 天的，无条件辞退。5\n",
      "124 ========== . 维护学校财产学校物品（包括发给员工使用的物品）均为学校财产，无论疏忽或有意损坏，当事人必须酌情赔偿。员工如凡有盗窃行为，学校将立即予解除劳动合同，并视情节轻重交由公安部门处理。企业形象是学校非常重要的财富，维护好企业形象是每个员工必须遵守的规则\n",
      "78 ========== 。员工必须严格遵守学校的企业文化，经营理念和管理制度，充分维护和支持学校的企业形象建设。任何人不得出现有损学校团队建设、诋毁学校企业管理和企业文化的语言行为\n",
      "76 ========== 。若教职工出现有损学校形象和荣誉的情况，由校长与当事人面谈，根据情节严重程度给与警告、处分，如不改正者，学校给予解除劳动合同处理，并保留法律追究的权利。\n",
      "481 ========== <table>\n",
      "<tr><td  >级别</td><td  >火车 （含高铁、动车、全列软席列车）</td><td  >轮船 （不包括 旅游船）</td><td  >飞机</td><td  >其他交通工具 （不包括出租 小汽车）</td></tr>\n",
      "<tr><td  >享受副部级 待遇及以上 人员</td><td  >火车软席（软座、软卧），高铁/动车商 务座，全列软席列车一等软座</td><td  >一等舱</td><td  >头等舱</td><td  >凭据报销</td></tr>\n",
      "<tr><td  >秘书长及副 秘书长</td><td  >火车软席（软座、软卧），高铁/动车一 等座，全列软席列车一等软座</td><td  >二等舱</td><td  >经济舱</td><td  >凭据报销</td></tr>\n",
      "<tr><td  >其余人员</td><td  >火车硬席（硬座、硬卧），高铁/动车二 等座、全列软席列车二等软座</td><td  >三等舱</td><td  >经济舱</td><td  >凭据报销</td></tr>\n",
      "</table>\n",
      "182 ========== <table><caption>伙食补助费参考以下标准：</caption>\n",
      "<tr><td  >地区</td><td  >伙食补助费标准</td></tr>\n",
      "<tr><td  >西藏、青海、新疆</td><td  >120 元/人、天</td></tr>\n",
      "<tr><td  >其他省份</td><td  >100 元/人、天</td></tr>\n",
      "</table>\n",
      "1296 ========== 差旅费用标准\n",
      "差旅费开支范围包括工作人员临时到常驻地以外地区公务出差所发生的城市间交通费、住宿费、伙食补助费和市内交通费。一、城市间交通费城市间交通费是指工作人员因公到常驻地以外地区出差乘坐火车、轮船、飞机等交通工具所发生的费用。1.出差人员在不影响公务、确保安全的前提下，选乘经济便捷的交通工具。2.出差人员要按照规定等级乘坐交通工具，未按规定乘坐的，超支部分自理。乘坐交通工具的等级见下表：\n",
      "<table>\n",
      "<tr><td  >级别</td><td  >火车 （含高铁、动车、全列软席列车）</td><td  >轮船 （不包括 旅游船）</td><td  >飞机</td><td  >其他交通工具 （不包括出租 小汽车）</td></tr>\n",
      "<tr><td  >享受副部级 待遇及以上 人员</td><td  >火车软席（软座、软卧），高铁/动车商 务座，全列软席列车一等软座</td><td  >一等舱</td><td  >头等舱</td><td  >凭据报销</td></tr>\n",
      "<tr><td  >秘书长及副 秘书长</td><td  >火车软席（软座、软卧），高铁/动车一 等座，全列软席列车一等软座</td><td  >二等舱</td><td  >经济舱</td><td  >凭据报销</td></tr>\n",
      "<tr><td  >其余人员</td><td  >火车硬席（硬座、硬卧），高铁/动车二 等座、全列软席列车二等软座</td><td  >三等舱</td><td  >经济舱</td><td  >凭据报销</td></tr>\n",
      "</table>\n",
      "备注：\n",
      "享受副部级待遇及以上人员出差，因工作需要，随行一人可乘坐同等级交通工具；乘坐飞机的，民航发展基金、燃油附加费可以凭据报销；乘坐飞机、火车、轮船等交通工具的，每人次可以购买交通意外保险一份。由我会统一购买交通意外保险的，不再重复购买。3.我会工作人员出差，高铁最短5 个小时内能够到达目的地的，原则上应乘坐高铁，如遇特殊情况，经分管副秘书长、秘书长审批后，可乘坐飞机。4.国内出差购买机票,应当按照厉行节约和支持本国航空公司发展的原则,优先购买通过政府采购方式确定的我国航空公司航班优惠机票。飞机票应通过定点供应商采购。二、住宿费住宿费是指工作人员因公出差期间入住宾馆（包括饭店、招待所，下同）发生的房租费用。享受副部级待遇及以上人员住普通套间，秘书长及以下人员住单间或标准间。出差人员应当在职务级别对应的住宿费标准限额内，选择安全、经济、便捷的宾馆住宿。国内出差发生的住宿，原则上应当按照住宿费标准执行，如遇特殊情况（如对方组织单位指定或推荐酒店等情况），超出费用标准的，须事先报分管资金的副秘书长、秘书长批准后方可报销。住宿费按照财政部《中央和国家机关工作人员赴地方差旅住宿费标准明细表》（财行[2016]71号）执行。三、伙食补助费伙食补助费是指对工作人员在因公出差期间给予的伙食补助费用。伙食补助费按出差自然(日历)天数计算，按规定标准包干使用。出差人员应当自行用餐。凡由接待单位统一安排用餐的，不再发放伙食补助费。\t\n",
      "184 ========== 2\n",
      "<table><caption>伙食补助费参考以下标准：</caption>\n",
      "<tr><td  >地区</td><td  >伙食补助费标准</td></tr>\n",
      "<tr><td  >西藏、青海、新疆</td><td  >120 元/人、天</td></tr>\n",
      "<tr><td  >其他省份</td><td  >100 元/人、天</td></tr>\n",
      "</table>\n",
      "101 ========== 四、市内交通费\n",
      "市内交通费是指工作人员因公出差期间发生的市内交通费用。出差人员的市内交通费按出差的自然（日历）天数计算，每人每天80 元包干使用。由接待单位或其他单位提供交通工具的，不再发放市内交通费。\n"
     ]
    }
   ],
   "source": [
    "for i in doc_data:\n",
    "    print(len(i), \"=\"*10,  i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c43187aa-d155-4296-b6d9-1c3ff4e5a243",
   "metadata": {},
   "outputs": [],
   "source": [
    "for excel_file_name in excel_files:\n",
    "    data = read_and_process_excel(excel_file_name)\n",
    "    df = pd.DataFrame(data[8:], columns=data[7])\n",
    "    data_excel = df.drop(columns=df.columns[11:17])\n",
    "    content = data_excel.to_markdown(index=False).replace(' ', \"\")\n",
    "    doc_data.append(f\"### 以下是中央和国家机关工作人员赴地方差旅住宿费标准明细表： \\n\\n {content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "a70d3edf-0de3-4db9-8dfc-cfbe0832fd06",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.documents import Document\n",
    "import uuid\n",
    "documents = []\n",
    "document_ids = []\n",
    "for idx, chunk in enumerate(doc_data):\n",
    "    is_table = 0\n",
    "    if \"table\" in chunk:\n",
    "        is_table = 1\n",
    "    if idx == len(doc_data) - 1:\n",
    "        is_table = 1\n",
    "    doc_id = str(uuid.uuid4())\n",
    "    document = Document(\n",
    "        page_content=chunk,\n",
    "        metadata={\"type\": \"ori\", \"is_table\": is_table})\n",
    "    documents.append(document)\n",
    "    document_ids.append(doc_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "0b50d33c-6e7a-4329-8060-77193bd03a8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_cls = RagEmbedding()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "84061f0c-e26e-4903-8b8b-a10e236f325c",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_txts = dict(zip(document_ids, documents))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "3309c584-83b2-46b9-bf96-43ceeb8d805b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open('./data/zhidu_db.pickl', 'wb') as file:\n",
    "    pickle.dump(doc_txts, file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "785bfd12-2c0b-4189-85a1-e7db7a36233b",
   "metadata": {},
   "outputs": [],
   "source": [
    "chroma_client = chromadb.HttpClient(host=\"localhost\", port=8000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a4d0618e-a773-4ac5-99b9-41b796efe269",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_db = Chroma.from_documents(documents,\n",
    "                                     embedding_cls.get_embedding_fun(),\n",
    "                                     client=chroma_client,\n",
    "                                     collection_name=\"zhidu_db\",\n",
    "                                     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "267ecce0-268b-4392-895a-726b9c210b4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = '迟到有什么规定？'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "856303d8-a6e5-4f81-a752-104f54e17ab9",
   "metadata": {},
   "outputs": [],
   "source": [
    "related_docs = embedding_db.similarity_search(query, k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c261effe-ab84-45d5-87be-d4bad644ee4e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(metadata={'is_table': 0, 'type': 'ori'}, page_content='。如遇紧急情况，口头申请请假的，应在上班后两天内办理补请假手续，未在规定时间内办理的，逾期无效，按旷工处理。4、无工作理由，超过上班时间到岗的，视为迟到；未到下班时间提前离校的，视为早退；中途未经批准离校，视为旷工；迟到、早退、旷工者按照相关办法处理'),\n",
       " Document(metadata={'is_table': 0, 'type': 'ori'}, page_content='第五节违规违纪处理\\n1. 迟到、早退(1)劳动考勤是公司支付薪资的依据和职工年度岗位考核内容之一。（2）迟到或早退60分钟以上（含60 分钟），每次视同缺勤1 天。（3）职工迟到、早退、脱岗累计超过3 次的（含），从第1 次起，每次扣减工资50 元。')]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "related_docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "64226aa3-a832-4619-ab0e-9be73f9597a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "chroma_client = chromadb.HttpClient(host=\"localhost\", port=8000)\n",
    "\n",
    "def delete_chroma_collection(name, embedding_fn, chroma_client):\n",
    "    collection = Chroma(name,\n",
    "                        embedding_fn,\n",
    "                        client=chroma_client)\n",
    "    Chroma.delete_collection(collection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fe53796-4499-44d3-b29d-11b3bfcb7dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_chroma_collection(\"zhidu_db\", embedding_cls.get_embedding_fun(), chroma_client)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "525d13cb-8c94-420a-99a0-9b45ecf88476",
   "metadata": {
    "tags": []
   },
   "source": [
    "## RAG问答流程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87555216-e8d9-44d4-a2bb-cb20a57e3975",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = RagLLM()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16c2b0ee-9c37-4a51-aeb3-b6e44dcce61e",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"\n",
    "你是企业员工助手，熟悉公司考勤和报销标准等规章制度，需要根据提供的上下文信息context来回答员工的提问。\\\n",
    "请直接回答问题，如果上下文信息context没有和问题相关的信息，请直接回答[不知道,请咨询HR] \\\n",
    "问题：{question} \n",
    "\"{context}\"\n",
    "回答：\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "856486bf-4d3a-46ca-b0ed-61fd7c3867ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_rag_pipline(querys, k=3):\n",
    "    for query in querys:\n",
    "        \n",
    "        related_docs = embedding_db.similarity_search(query, k=k)\n",
    "        context = \"\\n\".join([f\"上下文{i+1}: {doc.page_content} \\n\" \\\n",
    "                             for i, doc in enumerate(related_docs)])\n",
    "        print()\n",
    "        print()\n",
    "        print(\"#\"*100)\n",
    "        print(f\"query: {query}\")\n",
    "        print(f\"context: {context}\")\n",
    "        llm_prompt = prompt_template.replace(\"{question}\", query).replace(\"{context}\", context)\n",
    "        response = llm(llm_prompt, stream=True)\n",
    "        print(f\"response: \")\n",
    "        for chunk in response:\n",
    "            print(chunk.choices[0].text, end='', flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f580ea30-e20a-4747-b9cd-7679c8105a1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_rag_pipline(['加班有加班费吗？'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b288f7c-dcd4-4a0d-a3ff-65025ce14d9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_rag_pipline(['出差可以买意外保险吗？需要自己购买吗'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4833bc8f-df35-46f4-b1b3-f42b76ad2047",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_rag_pipline(['那个，我们公司有什么规定来着？ 您公司的具体规定需要参考您所在公司的员工手册或相关政策文件。通常，公司会有包括但不限于工作时间、请假流程、加班政策、行为准则、保密协议等方面的规章制度。建议查阅您的员工手册或向人力资源部门咨询以获取详细信息。'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "749936ec-40fe-4012-b92d-98705b8558f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_rag_pipline(['公司派我去上海，钱怎么弄？'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "440d68c5-3c9d-40c3-8abd-0571c2634953",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_rag_pipline(['如果我明天要出去，那个报销的事情怎么办？'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9616cf6-67d6-4c93-8332-31cb7b1ed7e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_rag_pipline(['如果我明天要出去，那个报销的事情怎么办？ 如果您明天需要外出，关于报销的事宜，请提前与您的直接上级或财务部门沟通。确保了解公司的报销流程和所需文件，如需预支费用，应按照公司规定提交申请。同时，保留好所有相关票据，以便返回后及时完成报销手续。如果可能，可以预先填写报销单的部分内容，以加快后续处理速度。'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d5fc4e-fe06-41da-a46c-41353a473db5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "language": "python",
   "name": "llm"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
